package com.ai.common.config;

import com.ai.common.cache.TimeExpiredPoolCache;
import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
@Slf4j
public class RateLimiter {
    //最大令牌数
    private final long maxTokens;
    private final AtomicLong tokens;
    //每秒补充的令牌数
    private final long refillTokensPerSecond;
    private long nextRefillTime;

    public RateLimiter(long maxTokens, long refillTokensPerSecond) {
        this.maxTokens = maxTokens;
        this.tokens = new AtomicLong(maxTokens);
        this.refillTokensPerSecond = refillTokensPerSecond;
        this.nextRefillTime = System.nanoTime();
    }

    /**
     * 尝试获取一个令牌
     * @return
     */
    public boolean tryAcquire() {
        refill();
        return tokens.getAndUpdate(current -> current > 0 ? current - 1 : 0) > 0;
    }

    private void refill() {
        long now = System.nanoTime();
        if (now < nextRefillTime) {
            return;
        }
        long newTokens = (now - nextRefillTime) * refillTokensPerSecond / TimeUnit.SECONDS.toNanos(1);
        if (newTokens > 0) {
            tokens.updateAndGet(current -> Math.min(current + newTokens, maxTokens));
            nextRefillTime = now;
        }
    }

    public static void main(String[] args) {
        String uid = "aaa";
        try {
            //获取用户令牌桶对象
            RateLimiter userRateLimiter = TimeExpiredPoolCache.getInstance().get(uid);
            if (userRateLimiter == null){
                //最大令牌数为10，每秒补充1个令牌的限流器
                userRateLimiter = new RateLimiter(3, 1);
                //保存用户令牌桶对象缓存 设置失效时间5分钟
                TimeExpiredPoolCache.getInstance().put(uid,userRateLimiter,5*60*1000L);
            }
            for (int i = 0; i < 50; i++) {
                //使用tryAcquire()方法进行限流,如果令牌桶中有令牌，则获取成功，执行接口调用；否则获取失败，返回接口访问频率过高的错误信息。在获取令牌时，限流器会自动补充令牌，确保令牌桶中的令牌数量不超过最大令牌数。
                if (userRateLimiter.tryAcquire()) {
                    // 执行接口调用
                    System.out.println(true);
                } else {
                    System.out.println(false);
                    // 返回接口访问频率过高的错误信息
                }
                try {
                    Thread.sleep(200); // 每0.2秒尝试一次访问接口
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        }catch (Exception e){
            log.error(e.getMessage(),e);
        }
    }
}
